Skip to content

Go HTTP 中间件:包住 Handler,也要守住响应边界 ​

日志、鉴权、CORS、限流看起来都能写成“处理前做一点事,处理后做一点事”。真正难的是它们叠在一起以后:鉴权失败有没有日志?跨域页面能不能读到 401?SSE 的 Flush 有没有被包装器藏掉?Handler 写了一半再 panic,还能不能返回 500?

这篇用标准 net/http 把这些关系拆开。所有实验在 Go 1.27.2、Linux/amd64 执行,完整中间件和测试放在文末。示例用于验证组合边界,不把一组短函数包装成开箱即用的生产安全组件。

一层包装的完整含义 ​

中间件接收一个 Handler,返回另一个 Handler:

go
func wrap(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // 进入内层之前
        next.ServeHTTP(w, r)
        // 内层返回之后
    })
}

调用 next 才会继续。若校验失败,写完错误后必须 return;只写 http.Error 并不会阻止下面的业务运行。组合 A(B(C(endpoint))) 的进入顺序是 A、B、C,正常返回顺序是 C、B、A。

正在绘制图示…

可切换适合窗口或原始尺寸
查看图源代码
flowchart LR
    A[Request ID] --> B[访问日志]
    B --> C[Recovery]
    C --> D[CORS 与预检]
    D --> E[入口限流]
    E --> F[身份验证]
    F --> G[用户配额与对象授权]
    G --> H[业务 Handler]
    H -. 状态和字节数逐层返回 .-> B
放大图示

原始尺寸下可横向或纵向滚动;按 Escape 关闭。

正在绘制图示…

查看图源代码
flowchart LR
    A[Request ID] --> B[访问日志]
    B --> C[Recovery]
    C --> D[CORS 与预检]
    D --> E[入口限流]
    E --> F[身份验证]
    F --> G[用户配额与对象授权]
    G --> H[业务 Handler]
    H -. 状态和字节数逐层返回 .-> B

这是一个可讨论的布局,不是唯一正确顺序。入口限流在鉴权前挡住粗暴流量;用户配额在鉴权后才能使用可信用户 ID。日志放外侧,才能记录预检、401 和 429。Recovery 能捕获它内层同一 goroutine 的 panic,却保护不了位于它外层的代码,因此最外侧自己的实现也要简单可靠。

中间件主要输入改变什么最重要的边界
请求 ID可信上游 ID 或新生成值Context、响应 Header不可信 Header 不能冒充可信身份
日志方法、路由、状态、耗时日志输出敏感数据与响应提交统计
recover同 goroutine 的 panic尚未提交时的错误响应不回滚业务,也不拦截子 goroutine
CORSOrigin、预检信息浏览器可读响应的许可不是身份验证
鉴权凭据与已验证身份是否继续登录不等于有对象操作权限
限流时间、额度、身份维度是否拒绝或等待请求频率不等于并发量
超时deadline、取消信号停止等待或限制响应无法强停忽略取消的业务
压缩协商结果、内容类型响应字节与 Header缓冲、Flush、长度与缓存变化

请求 ID 和日志:先能找到一次请求,再解释发生了什么 ​

请求 ID 是排查线索。入口创建 ID,写进响应 Header,并放入 Context 传给下游日志。Context key 用包内私有类型,避免不同组件碰巧使用同一个字符串 key。context.WithValue

如果接收来自外部的 X-Request-ID,应限制长度、字符集和来源。也可以总是生成内部 ID,把经过校验的外部 ID 另存一列。它不是访问令牌,不能凭一个 ID 授权。文末用原子递增的 local-1 演示传播,仅保证本进程序列;真实多实例系统要使用合适的唯一标识或追踪协议。W3C Trace Context

日志至少能回答方法、路由、结果、耗时和关联 ID。记录路由模板 /users/{id} 有助于控制指标基数;完整 URL 可能带搜索词、令牌或其他敏感值,不应默认原样输出。记录 Header 和 Body 更要按白名单及脱敏策略处理。

状态码为什么需要包 ResponseWriter ​

ResponseWriter 没有“取最后状态码”的方法。日志包装器要拦截 WriteHeader 与 Write:第一次普通最终状态生效;如果先 Write,状态就是隐式 200。103 等临时响应不应占用最终状态位置;101 Switching Protocols 是协议切换,Go 将其作为终结状态处理,包装器也必须标记为已提交(Go 1.27.2 server.go 的 101 特例)。字节数应累加实际返回的 n,不要无条件加 len(p)。ResponseWriter

本次测试写入 hello,去掉时间字段以保持稳定,输出:

json
{"level":"INFO","msg":"request","method":"GET","status":200,"bytes":5}

这里的 5 是本层 Writer 接受的字节,不是客户端实际收到的字节。日志层放在压缩前后,统计的还可能分别是原始体积和压缩体积。另外的状态测试确认 103 后仍可记录最终 200,而 101 后不再改写为 500。生产日志还要单独记录 panic、中断、超时与写错误;不能把“尚未写出任何状态时的默认 200”误标成成功。文末日志函数只演示正常路径的状态与字节统计,实际部署应让它包住 Recovery 并补上中断字段。

ResponseWriter 包装器会改变接口能力 ​

下面看似只包了一层:

go
wrapped := &struct{ http.ResponseWriter }{w}

它只暴露嵌入接口声明的 Header、Write、WriteHeader,不会自动把底层的 Flush、Hijack 等附加能力带出来。本次负例中,对它调用 ResponseController.Flush 得到 http.ErrNotSupported。

Go 当前提供 ResponseController,它可以寻找 FlushError、Flush,也会沿 Unwrap() http.ResponseWriter 继续找底层。一个计数包装器可以实现:

go
func (w *measuredWriter) Unwrap() http.ResponseWriter {
    return w.ResponseWriter
}
func (w *measuredWriter) FlushError() error {
    if w.status == 0 {
        w.WriteHeader(http.StatusOK)
    }
    return http.NewResponseController(w.ResponseWriter).Flush()
}

本次测试确认 Recorder 被标记为 Flushed,日志仍能记录隐式提交的 200。为什么还要写 FlushError,不能只写 Unwrap?因为 Flush 可能提交响应,计数器和 Recovery 必须知道这一点;如果直接绕到底层,包装器的“尚未提交”标记就可能失真。ResponseController 源码,Go 1.27.2

Unwrap 也不会让旧代码的 w.(http.Flusher) 自动成功。调用方若直接断言旧接口,需要包装器相应支持;Hijacker、Pusher、io.ReaderFrom 等同样要逐个评估,不能伪造一个底层不支持的能力。WebSocket 升级、SSE、文件传输都要测试经过完整链后的行为。

recover:响应写出去前后是两种处理 ​

标准 HTTP Server 本身会 recover Handler 的 panic,但通常是关闭连接或重置 HTTP/2 流,不是替应用返回整齐的 JSON 500。应用 Recovery 的用途,是在安全范围内记录受控日志,并在响应尚未提交时给出统一错误。Handler 文档

panic 发生时可做的事不能假设什么
尚未提交最终响应记录堆栈,返回受控 500数据库事务已自动回滚
已提交部分响应按已有流协议处理,或中断连接/流后写的 500 会替换之前的 200
新启动的 goroutine 内在该 goroutine 自己的边界处理外层 HTTP Recovery 能跨 goroutine 接住

文末 recoverBeforeCommit 在已提交后重新 panic http.ErrAbortHandler,交给 Server 中断响应,不把内部堆栈写给客户端。测试分别断言提交前得到 500、提交后继续抛出这个中断标记。它还保留已有的 ErrAbortHandler,不把明确的中断请求转换成成功或普通错误页。

Recovery 不能修复程序已经破坏的共享状态,也不能恢复已发出的邮件或已提交的付款。事务与副作用的原子性,应在业务层设计;HTTP 边界只能诚实结束这次请求。

CORS 和鉴权处理的是两个问题 ​

CORS 由浏览器执行,用来决定跨源脚本能否读取响应,并在某些请求前先执行预检。命令行程序、另一台服务器或攻击者自己的客户端不需要遵守浏览器的 CORS 限制。Origin 允许值匹配也不证明请求者是谁。Fetch CORS 协议

本次链允许 https://app.example,鉴权需要公开的测试令牌。实际断言:

请求结果原因
允许的 Origin,无令牌401CORS 许可没有替代鉴权
不带 Origin,有测试令牌204非浏览器请求仍按凭据处理
非允许 Origin,有测试令牌403本示例主动拒绝该 Origin
允许 Origin 的合法 OPTIONS 预检,无令牌204只回答允许的方法/头,不执行业务

预检通常不携带实际请求的用户凭据,所以让 CORS 在常规鉴权前处理符合本策略的预检,否则浏览器可能连真正请求都发不出。实际业务请求仍要鉴权。CORS 错误只是浏览器看不到响应的原因之一;某些跨源请求能先发出去,再因响应不满足 CORS 而不可读,所以不能用 CORS 代替 CSRF 防护。

允许凭据时,Access-Control-Allow-Origin 不能用 *,必须按允许列表回应具体 Origin,并正确设置 Access-Control-Allow-Credentials。动态回应 Origin 时加 Vary: Origin,否则共享缓存可能把一种 Origin 的许可误用给另一种。预检还要检查方法与 Header 列表,Header 名按大小写不敏感处理。文末实现仅演示单一 Origin、GET 和 Authorization,不支持通配策略、Cookie 凭据或完整 CORS 配置系统。

鉴权本身应分两步:验证凭据得到可信主体,再检查主体对操作和对象的权限。JWT 能被解码不代表签名、算法、发行者、受众和时效都通过(见 RFC 8725 §3);从 Context 拿到 user ID 也不代表能读取任意 /users/{id}。公开常量 Bearer test-only 只是测试开关,绝不能当生产认证方案。

限流先说清限制的是什么 ​

机制限制对象典型用途
令牌桶一段时间内允许的速率,并容忍有限突发API 请求频率
固定/滑动窗口时间窗口里的次数用户每日额度、短期计数
并发信号量正在执行的数量昂贵查询、图片处理
有界队列等待处理的任务量把过载显式暴露出来

令牌桶有容量和补充速度。容量 2、每秒补充 1 个令牌,同一时刻请求三次,再过一秒请求一次:

go
now := time.Unix(0, 0)
b := bucket{tokens: 2, capacity: 2, perSecond: 1, last: now}
fmt.Println(b.allow(now), b.allow(now), b.allow(now))
fmt.Println(b.allow(now.Add(time.Second)))

确定性时钟实验得到:

text
true true false
true

第三次被拒绝,因为桶已经空了。它不会阻止前两次每个运行一分钟;要限制同时运行数量,还需并发控制。官方 golang.org/x/time/rate 提供令牌桶实现;本文自己写小模型,是为了展示数量与时间的关系,没有替代成熟库的意图。rate 包

选择 IP 限流时要小心共享 NAT,也不要盲信任客户端填的 X-Forwarded-For。按用户限流要先获得可信身份。每个键都建一个永不过期的桶,会被大量不同 key 撑爆内存;多实例各自限额会让总配额随实例数变化。全局配额需要额外协调,并明确存储故障时放行还是拒绝。

拒绝可以返回 429;能给出可靠等待时间时再发 Retry-After。示例没有凭空填写这个 Header。等待型限流还要接受 Context 取消,并限制队列长度,避免把过载藏成越来越多的 goroutine。RFC 6585 §4

超时取消业务,和超时结束响应,要分别设计 ​

context.WithTimeout 到期会关闭 Done。它不会把 Handler 从 CPU 上摘下来,也不会自动往网络写 504。本次测试甚至让已取消 Context 的 Handler 继续写,决定结果的处理函数只有两行:

go
<-r.Context().Done()
io.WriteString(w, "still ran") // 仍然主动写入,取消并没有接管响应

结果仍是:

text
canceled context without response policy: 200 "still ran"

这不是推荐写法,而是说明必须合作式停止。下游 HTTP、数据库调用要传入 Context;循环和等待要检查 Done;停止之后该回什么错误,由 Handler 或明确的响应策略决定。context 文档

标准库 http.TimeoutHandler 则管理响应:正常响应暂存,超时后返回 503,之后内层写入得到 ErrHandlerTimeout。本次内层故意忽略取消,等外层已经超时后再写:

go
h := http.TimeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    <-release
    _, err := io.WriteString(w, "too late")
    written <- err
}), 10*time.Millisecond, "timed out")
h.ServeHTTP(recorder, request) // 到期返回,此时内层仍等 release
close(release)
err := <-written

其中 release 和 written 是测试创建的 channel。实测:

text
TimeoutHandler: 503; late write rejected=true

它没有杀掉内层,内层直到测试释放 channel 才停止。TimeoutHandler 不支持 Flusher、Hijacker,因此不适合直接包 SSE 或 WebSocket;缓冲响应本身也有内存代价。TimeoutHandler

需求主要手段剩下的责任
慢请求头/慢上传/慢写Server I/O 超时与大小限制按上传、短请求、流式请求区分预算
停止下游工作Context deadline 与传播下游和业务合作退出
到期后给固定响应TimeoutHandler 或谨慎设计的响应缓冲器已提交响应不能改写,内层可能仍在运行
SSE 长时间连接每次写的 deadline、心跳、取消不用短的整次响应超时误杀正常流

不要自己开一个 goroutine 写真实 ResponseWriter,外层超时后又同时写错误。两个 goroutine 竞争同一响应,且内层可能在外层返回后继续使用 Writer,已经违反 Handler 生命周期。

压缩减少字节,也引入一层状态 ​

gzip 适合可压缩且足够大的文本。JPEG、ZIP 等已经压缩的数据可能收益很小;极小响应可能因 Header 和压缩结构反而变大。是否启用应该根据媒体类型、大小、CPU 预算和客户端 Accept-Encoding 协商,而不是所有响应无条件包 gzip。

协商后设置 Content-Encoding: gzip,添加 Vary: Accept-Encoding,不要继续保留原始 Body 的 Content-Length。涉及 ETag、Range、HEAD、204/304、已有 Content-Encoding 的路径也需要明确策略。Accept-Encoding 的 q=0 表示不接受,不能只用字符串包含 gzip 判断。RFC 9110 §8.4、§12.5.3

流式压缩还要按顺序刷两层:

go
if err := z.Flush(); err != nil {
    return err
}
return http.NewResponseController(underlying).Flush()

这里 z 是 *gzip.Writer,underlying 是它写入的底层 http.ResponseWriter。先让压缩器吐出已有数据,再让 HTTP 层发送。结束时必须 Close 压缩器,写完尾部信息;Close 不会替你关闭底层连接。本次测试解压得到 hello 并读到正常结束,同时验证 HTTP Flush 被调用。gzip.Writer.Flush / Close

文末只测试“已经决定使用 gzip 后如何写、Flush、Close”,没有实现完整协商中间件。实际系统通常对 SSE 单独禁用压缩或验证延迟效果,也要评估把攻击者可控文本与秘密共同压缩造成的长度侧信道风险,机制背景见 RFC 7457 §2.6。

资料与实测边界 ​

go test、race 和 vet 全通过。已验证拒绝路径不执行业务、CORS 与鉴权独立、Recovery 提交边界、Flush 包装反例、令牌数量、Context 与响应超时差别、gzip 两层刷新与尾部。没有运行真实浏览器 CORS、集群限流、生产身份系统或压缩安全攻击实验。

完整代码与负例 ​

把以下文件放在同一目录,使用 Go 1.27.2:

sh
go test -v -count=1 ./...
go test -race -count=1 ./...
go vet ./...

go.mod ​

go.mod
module example.com/webapi/04_middleware

go 1.27.2

middleware_test.go ​

代码块 · 155 行 · 点击收起 / 展开
go
package middlewaredemo

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
)

type requestKey struct{}

var sequence atomic.Uint64

func requestID(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		id := fmt.Sprintf("local-%d", sequence.Add(1))
		w.Header().Set("X-Request-ID", id)
		next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), requestKey{}, id)))
	})
}
func cors(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		origin := r.Header.Get("Origin")
		w.Header().Add("Vary", "Origin")
		if origin != "" && origin != "https://app.example" {
			http.Error(w, "origin rejected", 403)
			return
		}
		if origin != "" {
			w.Header().Set("Access-Control-Allow-Origin", origin)
		}
		if r.Method == "OPTIONS" && r.Header.Get("Access-Control-Request-Method") != "" {
			w.Header().Add("Vary", "Access-Control-Request-Method")
			w.Header().Add("Vary", "Access-Control-Request-Headers")
			if origin == "" || r.Header.Get("Access-Control-Request-Method") != "GET" {
				http.Error(w, "preflight rejected", 403)
				return
			}
			for _, h := range strings.Split(r.Header.Get("Access-Control-Request-Headers"), ",") {
				h = strings.TrimSpace(h)
				if h != "" && !strings.EqualFold(h, "Authorization") {
					http.Error(w, "headers rejected", 403)
					return
				}
			}
			w.Header().Set("Access-Control-Allow-Methods", "GET")
			w.Header().Set("Access-Control-Allow-Headers", "Authorization")
			w.WriteHeader(204)
			return
		}
		next.ServeHTTP(w, r)
	})
}
func auth(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// A fixed, public test token is not a production authentication system.
		if r.Header.Get("Authorization") != "Bearer test-only" {
			w.Header().Set("WWW-Authenticate", `Bearer realm="example"`)
			http.Error(w, "authentication required", 401)
			return
		}
		next.ServeHTTP(w, r)
	})
}

type committedWriter struct {
	http.ResponseWriter
	committed bool
}

func (w *committedWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
func (w *committedWriter) FlushError() error {
	if !w.committed {
		w.WriteHeader(200)
	}
	return http.NewResponseController(w.ResponseWriter).Flush()
}

func (w *committedWriter) WriteHeader(code int) {
	if code >= 200 || code == http.StatusSwitchingProtocols {
		w.committed = true
	}
	w.ResponseWriter.WriteHeader(code)
}
func (w *committedWriter) Write(p []byte) (int, error) {
	w.committed = true
	return w.ResponseWriter.Write(p)
}
func recoverBeforeCommit(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		cw := &committedWriter{ResponseWriter: w}
		defer func() {
			if p := recover(); p != nil {
				if p == http.ErrAbortHandler || cw.committed {
					panic(http.ErrAbortHandler)
				}
				// A real server logs the stack to a restricted sink, never to this response.
				http.Error(cw, "internal error", 500)
			}
		}()
		next.ServeHTTP(cw, r)
	})
}
func TestChain(t *testing.T) {
	var calls atomic.Int32
	endpoint := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls.Add(1)
		if r.Context().Value(requestKey{}) == nil {
			panic("missing request id")
		}
		w.WriteHeader(204)
	})
	h := requestID(cors(recoverBeforeCommit(auth(endpoint))))
	for _, tc := range []struct {
		origin, token string
		want          int
	}{
		{"https://app.example", "", 401}, {"", "Bearer test-only", 204},
		{"https://evil.example", "Bearer test-only", 403},
	} {
		r := httptest.NewRequest("GET", "/", nil)
		r.Header.Set("Origin", tc.origin)
		r.Header.Set("Authorization", tc.token)
		w := httptest.NewRecorder()
		h.ServeHTTP(w, r)
		if w.Code != tc.want || w.Header().Get("X-Request-ID") == "" {
			t.Fatalf("%d %v", w.Code, w.Header())
		}
	}
	r := httptest.NewRequest("OPTIONS", "/", nil)
	r.Header.Set("Origin", "https://app.example")
	r.Header.Set("Access-Control-Request-Method", "GET")
	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)
	if w.Code != 204 || calls.Load() != 1 {
		t.Fatalf("status=%d endpoint calls=%d", w.Code, calls.Load())
	}
	t.Log("CORS allowed origin still needs auth; Origin-free client still needs auth; preflight does not call business handler")
}
func TestRecoveryBoundary(t *testing.T) {
	w := httptest.NewRecorder()
	recoverBeforeCommit(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { panic("boom") })).ServeHTTP(w, httptest.NewRequest("GET", "/", nil))
	if w.Code != 500 {
		t.Fatal(w.Code)
	}
	defer func() {
		if p := recover(); p != http.ErrAbortHandler {
			t.Fatalf("post-commit panic=%v", p)
		}
	}()
	recoverBeforeCommit(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte("partial")); panic("boom") })).ServeHTTP(httptest.NewRecorder(), httptest.NewRequest("GET", "/", nil))
}

teaching_test.go ​

代码块 · 255 行 · 点击收起 / 展开
go
package middlewaredemo

import (
	"bytes"
	"compress/gzip"
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"
)

type measuredWriter struct {
	http.ResponseWriter
	status, bytes int
}

func (w *measuredWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
func (w *measuredWriter) WriteHeader(code int) {
	if code >= 100 && code < 200 && code != http.StatusSwitchingProtocols {
		w.ResponseWriter.WriteHeader(code)
		return
	}
	if w.status != 0 {
		return
	}
	w.status = code
	w.ResponseWriter.WriteHeader(code)
}
func (w *measuredWriter) Write(p []byte) (int, error) {
	if w.status == 0 {
		w.WriteHeader(200)
	}
	n, e := w.ResponseWriter.Write(p)
	w.bytes += n
	return n, e
}
func (w *measuredWriter) FlushError() error {
	if w.status == 0 {
		w.WriteHeader(200)
	}
	return http.NewResponseController(w.ResponseWriter).Flush()
}
func accessLog(logger *slog.Logger, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mw := &measuredWriter{ResponseWriter: w}
		defer func() {
			status := mw.status
			if status == 0 {
				status = 200
			}
			logger.InfoContext(r.Context(), "request", "method", r.Method, "status", status, "bytes", mw.bytes)
		}()
		next.ServeHTTP(mw, r)
	})
}
func TestLoggingAndFlush(t *testing.T) {
	var out bytes.Buffer
	logger := slog.New(slog.NewJSONHandler(&out, &slog.HandlerOptions{ReplaceAttr: func(groups []string, a slog.Attr) slog.Attr {
		if a.Key == slog.TimeKey {
			return slog.Attr{}
		}
		return a
	}}))
	h := accessLog(logger, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "hello")
		if err := http.NewResponseController(w).Flush(); err != nil {
			t.Error(err)
		}
	}))
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest("GET", "/", nil))
	if !rec.Flushed {
		t.Fatal("not flushed")
	}
	t.Log(out.String())
	var plain http.ResponseWriter = &struct{ http.ResponseWriter }{rec}
	if err := http.NewResponseController(plain).Flush(); !errors.Is(err, http.ErrNotSupported) {
		t.Fatal(err)
	}
	t.Log("plain wrapper loses Flush; explicit FlushError keeps it")
}

type bucket struct {
	mu                          sync.Mutex
	tokens, capacity, perSecond float64
	last                        time.Time
}

func (b *bucket) allow(now time.Time) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	elapsed := now.Sub(b.last).Seconds()
	if elapsed > 0 {
		b.tokens = min(b.capacity, b.tokens+elapsed*b.perSecond)
		b.last = now
	}
	if b.tokens < 1 {
		return false
	}
	b.tokens--
	return true
}
func Example_bucket() {
	now := time.Unix(0, 0)
	b := bucket{tokens: 2, capacity: 2, perSecond: 1, last: now}
	fmt.Println(b.allow(now), b.allow(now), b.allow(now))
	fmt.Println(b.allow(now.Add(time.Second)))
	// Output:
	// true true false
	// true
}
func limited(b *bucket, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !b.allow(time.Now()) {
			http.Error(w, "rate limited", 429)
			return
		}
		next.ServeHTTP(w, r)
	})
}
func TestLimit(t *testing.T) {
	b := &bucket{tokens: 1, capacity: 1, perSecond: 0, last: time.Now()}
	h := limited(b, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(204) }))
	for _, want := range []int{204, 429} {
		w := httptest.NewRecorder()
		h.ServeHTTP(w, httptest.NewRequest("GET", "/", nil))
		if w.Code != want {
			t.Fatal(w.Code)
		}
	}
}

func TestContextAndResponseTimeout(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	r := httptest.NewRequest("GET", "/", nil).WithContext(ctx)
	w := httptest.NewRecorder()
	http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { <-r.Context().Done(); io.WriteString(w, "still ran") }).ServeHTTP(w, r)
	t.Logf("canceled context without response policy: %d %q", w.Code, w.Body.String())
	release := make(chan struct{})
	written := make(chan error, 1)
	h := http.TimeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		<-release
		_, err := io.WriteString(w, "too late")
		written <- err
	}), 10*time.Millisecond, "timed out")
	w = httptest.NewRecorder()
	h.ServeHTTP(w, httptest.NewRequest("GET", "/", nil))
	if w.Code != 503 {
		t.Fatal(w.Code)
	}
	close(release)
	err := <-written
	if !errors.Is(err, http.ErrHandlerTimeout) {
		t.Fatal(err)
	}
	t.Logf("TimeoutHandler: %d; late write rejected=%v", w.Code, errors.Is(err, http.ErrHandlerTimeout))
}

type gzipWriter struct {
	http.ResponseWriter
	z *gzip.Writer
}

func (w *gzipWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
func (w *gzipWriter) Write(p []byte) (int, error) { return w.z.Write(p) }
func (w *gzipWriter) FlushError() error {
	if err := w.z.Flush(); err != nil {
		return err
	}
	return http.NewResponseController(w.ResponseWriter).Flush()
}
func TestGzipFlushAndClose(t *testing.T) {
	rec := httptest.NewRecorder()
	rec.Header().Set("Content-Encoding", "gzip")
	rec.Header().Add("Vary", "Accept-Encoding")
	z := gzip.NewWriter(rec)
	w := &gzipWriter{ResponseWriter: rec, z: z}
	if _, err := io.WriteString(w, "hello"); err != nil {
		t.Fatal(err)
	}
	if err := http.NewResponseController(w).Flush(); err != nil {
		t.Fatal(err)
	}
	if !rec.Flushed {
		t.Fatal("not flushed")
	}
	if err := z.Close(); err != nil {
		t.Fatal(err)
	}
	r, err := gzip.NewReader(bytes.NewReader(rec.Body.Bytes()))
	if err != nil {
		t.Fatal(err)
	}
	defer r.Close()
	body, err := io.ReadAll(r)
	if err != nil || string(body) != "hello" {
		t.Fatalf("%q %v", body, err)
	}
	t.Log("gzip inner flush, HTTP flush, and final trailer verified")
}

// Header sink deliberately does not model network framing; these assertions
// test the wrappers' state transitions, including 101's terminal status.
type headerSink struct {
	h     http.Header
	codes []int
}

func (w *headerSink) Header() http.Header {
	if w.h == nil {
		w.h = make(http.Header)
	}
	return w.h
}
func (w *headerSink) WriteHeader(code int)        { w.codes = append(w.codes, code) }
func (w *headerSink) Write(p []byte) (int, error) { return len(p), nil }
func TestInformationalAndUpgradeState(t *testing.T) {
	mw := &measuredWriter{ResponseWriter: new(headerSink)}
	mw.WriteHeader(103)
	if mw.status != 0 {
		t.Fatal(mw.status)
	}
	mw.WriteHeader(200)
	if mw.status != 200 {
		t.Fatal(mw.status)
	}
	mw = &measuredWriter{ResponseWriter: new(headerSink)}
	mw.WriteHeader(101)
	mw.WriteHeader(500)
	if mw.status != 101 {
		t.Fatal(mw.status)
	}
	cw := &committedWriter{ResponseWriter: new(headerSink)}
	cw.WriteHeader(103)
	if cw.committed {
		t.Fatal("103 committed")
	}
	cw.WriteHeader(200)
	if !cw.committed {
		t.Fatal("200 not committed")
	}
	cw = &committedWriter{ResponseWriter: new(headerSink)}
	cw.WriteHeader(101)
	if !cw.committed {
		t.Fatal("101 not committed")
	}
	t.Log("103 then 200 tracked; 101 treated as terminal")
}

保留各书来源与许可 · 示例运行条件以正文为准